import json
import logging
import os
import sys
from datetime import datetime
from typing import Any, Dict, List

from Bio import Entrez

from tools.utils import BioInfoError, ConfigManager, DataFormatter

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)


class PubMedManager:
    """PubMed API 管理器"""
    
    def __init__(self, config_manager: ConfigManager):
        """初始化PubMed管理器"""
        self.config = config_manager
        # 设置邮箱，这是NCBI要求的
        Entrez.email = config_manager.get('NCBI', 'email')
    
    def search_pubmed(self, query: str, limit: int = None) -> List[str]:
        """使用Entrez.esearch进行PubMed搜索，返回文章ID列表。
        
        Args:
            query: 搜索查询词
            limit: 返回结果数量限制，如果为None则使用配置文件中的值
        """
        try:
            # 确定返回结果数量
            retmax = limit if limit is not None else int(self.config.get('NCBI', 'retmax'))
            
            # 使用Entrez.esearch进行搜索
            handle = Entrez.esearch(
                db=self.config.get('NCBI', 'db'),
                term=query,
                sort=self.config.get('NCBI', 'sort'),
                retmax=retmax
            )
            record = Entrez.read(handle)
            handle.close()
            # 检查 record 是否为字典类型且包含 'IdList' 键
            if isinstance(record, dict) and 'IdList' in record:
                id_list = record['IdList']
            else:
                id_list = []
            logger.info(f"找到 {len(id_list)} 条匹配的PubMed记录")
            return id_list
        except Exception as e:
            raise BioInfoError(f"PubMed搜索失败: {e}")
    
    def fetch_pubmed_records(self, id_list: List[str]) -> Dict[str, Any]:
        """使用Entrez.efetch获取文章记录。"""
        try:
            # 使用Entrez.efetch获取文章记录
            handle = Entrez.efetch(
                db=self.config.get('NCBI', 'db'),
                id=id_list,
                retmode=self.config.get('NCBI', 'retmode')
            )
            records = Entrez.read(handle)
            handle.close()
            return records
        except Exception as e:
            raise BioInfoError(f"获取PubMed记录失败: {e}")
    
    def extract_info(self, records: Dict[str, Any]) -> List[Dict[str, str]]:
        """从文章记录中提取文章信息。"""
        results = []
        try:
            for record in records["PubmedArticle"]:
                article = record["MedlineCitation"]["Article"]
                # 提取文章标题
                title = article["ArticleTitle"]
                # 提取文章摘要，如果有摘要的话
                abstract = article.get("Abstract", {}).get("AbstractText", ["No abstract available"])
                abstract = " ".join(abstract)
                # 提取文章出处
                journal = article["Journal"]["Title"]
                pub_date = article["Journal"]["JournalIssue"]["PubDate"]
                year = pub_date.get("Year", "Unknown")
                volume = article["Journal"]["JournalIssue"].get("Volume", "Unknown")
                issue = article["Journal"]["JournalIssue"].get("Issue", "Unknown")
                pages = article.get("Pagination", {}).get("MedlinePgn", "Unknown")
                citation = f"{journal}. {year};{volume}({issue}):{pages}"

                results.append({
                    "title": title,
                    "abstract": abstract,
                    "citation": citation
                })
            logger.info(f"成功提取 {len(results)} 条文章信息")
            return results
        except Exception as e:
            raise BioInfoError(f"提取文章信息失败: {e}")


class PubMedDataFormatter:
    """PubMed特定的数据格式化工具"""
    
    @staticmethod
    def convert_to_markdown(data: List[Dict[str, str]], filename: str, 
                          query_type: str = 'pubmed') -> None:
        """将PubMed数据转换为Markdown格式"""
        md_filename = filename.replace('.json', '.md')
        
        try:
            with open(md_filename, 'w', encoding='utf-8') as md_file:
                # 写入PubMed查询结果的Markdown格式
                for i, result in enumerate(data, 1):
                    md_file.write(f"### 结果 {i}\n")
                    md_file.write(f"**标题:**\n**{result['title']}**\n")
                    md_file.write(f"**摘要:**\n{result['abstract']}\n")
                    md_file.write(f"**出处:**\n{result['citation']}\n")
                    md_file.write("\n")
            
            logger.info(f"Markdown结果已保存到: {md_filename}")
        except Exception as e:
            raise BioInfoError(f"保存Markdown文件失败: {e}")


def pubmed_search(args):
    """PubMed搜索功能"""
    try:
        config_manager = ConfigManager(args.config)
        pubmed_manager = PubMedManager(config_manager)
        
        # 搜索文章
        logger.info(f"开始PubMed搜索: {args.query}, 限制结果数: {args.limit}")
        id_list = pubmed_manager.search_pubmed(args.query, args.limit)
        
        if not id_list:
            print("未找到与查询关键词匹配的文章，请尝试其他关键词。")
            return
        
        # 获取文章记录
        records = pubmed_manager.fetch_pubmed_records(id_list)
        # 提取文章信息
        results = pubmed_manager.extract_info(records)

        # 创建保存目录
        result_dir = os.path.join('result', 'pubmed')
        os.makedirs(result_dir, exist_ok=True)
        
        # 获取当前日期
        current_date = datetime.now().strftime("%Y%m%d")
        # 生成文件名，替换可能导致文件系统问题的字符
        safe_query = args.query.replace('/', '_').replace('\\', '_').replace(':', '_')
        filename = os.path.join(
            result_dir, 
            f"{safe_query}_size{args.limit}_{current_date}.json"
        )
        
        # 保存结果
        DataFormatter.save_json(results, filename)
        
        # 转换为Markdown格式并保存
        PubMedDataFormatter.convert_to_markdown(results, filename)
        
        # 创建ZIP归档（包含JSON和Markdown文件）
        if args.format == 'all':
            DataFormatter.create_zip_archive(filename, include_csv=False, include_md=True)
        
        print(f"搜索完成！结果已保存到 {filename} 和 {filename.replace('.json', '.md')}")
    except BioInfoError as e:
        logger.error(f"错误: {e}")
        print(f"错误: {e}")
        sys.exit(1)